`include "defines.v"

module ID_Stage (
    input  wire [`BUSLEN-1:0]           pc_i,
    output wire [`BUSLEN-1:0]           pc_o,
    input  wire                         pc_valid_i,
    output wire                         pc_valid_o,
    input  wire [`INSTLEN-1:0]          inst,

    // ALU Oprends
    input  wire                         forward_rs1_flag,
    input  wire                         forward_rs2_flag,
    input  wire                         forward_csr_flag,
    input  wire [`XLEN-1:0]             forward_rs1_data,
    input  wire [`XLEN-1:0]             forward_rs2_data, 
    input  wire [`XLEN-1:0]             forward_csr_data,

    input  wire [`XLEN-1:0]             rs1_data_i,
    input  wire [`XLEN-1:0]             rs2_data_i,
    input  wire [`XLEN-1:0]             csr_data_i,

    output reg  [`XLEN-1:0]             regbus_A,
    output wire [`XLEN-1:0]             regbus_B,
    output reg  [`XLEN-1:0]             imm,

    // ALU Operation Control
    output wire [ 1: 0]                 ALU_oprend_A_src,
    output wire [ 1: 0]                 ALU_oprend_B_src,
    output reg  [ 2: 0]                 ALU_op,
    output wire [ 2: 0]                 ALU_out_ext_type,

    // Regfile/CSR R/W
    output wire                         rs1_r_en,
    output wire [`REGFILE_ADDR_LEN-1:0] rs1_r_addr,
    output wire                         rs2_r_en,
    output wire [`REGFILE_ADDR_LEN-1:0] rs2_r_addr,

    output wire                         rd_w_en,
    output wire                         rd_w_src,
    output wire [`REGFILE_ADDR_LEN-1:0] rd_w_addr,

    output wire                         csr_r_en,
    output wire                         csr_w_en,
    output wire [`CSR_ADDR_LEN-1:0]     csr_rw_addr,

    // Memory Access
    output wire                         mem_r_en,
    output wire [ 2: 0]                 mem_r_ext_type,
    output wire                         mem_w_en,
    output wire [ 1: 0]                 mem_w_size,

    // Jump/Branch Signals
    output wire                         adder_src,
    output wire                         inst_jump,
    output wire [ 2: 0]                 inst_branch_type,

    // Privileged Instruction / Exception
    output wire                         is_system_inst,
    output wire                         inst_ecall,
    output wire                         inst_mret,
    output wire                         inst_fencei,
    output wire                         inst_wfi
);

    wire [ 4: 0] opcode;
    wire [ 2: 0] funct3;
    wire [ 6: 0] funct7;
    wire [11: 0] funct12;
    wire [ 5: 0] shamt;

    assign opcode      = inst[ 6: 2];
    assign rd_w_addr   = inst[11: 7];
    assign funct3      = inst[14:12];
    assign rs1_r_addr  = inst[19:15];
    assign rs2_r_addr  = inst[24:20];
    assign funct7      = inst[31:25];
    assign shamt       = inst[25:20];
    assign csr_rw_addr = inst[31:20];
    assign funct12     = inst[31:20];

    //----------Regbus Mux----------//
    // Forwarding
    wire [`XLEN-1:0] rs1_data;
    wire [`XLEN-1:0] rs2_data;
    wire [`XLEN-1:0] csr_data;

    assign rs1_data = (forward_rs1_flag & rs1_r_en)       ?  forward_rs1_data  :  rs1_data_i  ;
    assign rs2_data = (forward_rs2_flag & rs2_r_en)       ?  forward_rs2_data  :  rs2_data_i  ;
    assign csr_data = (forward_csr_flag & is_system_inst) ?  forward_csr_data  :  csr_data_i  ;

    // Regbus A
    wire [1:0] regbusA_src;

    assign regbusA_src[1] = is_system_inst & funct3[1] & funct3[0];
    assign regbusA_src[0] = is_system_inst & funct3[2];

    always@(*) begin
        case(regbusA_src)
            2'b00   : regbus_A = rs1_data;
            2'b01   : regbus_A = {59'b0,rs1_r_addr};
            2'b10   : regbus_A = ~rs1_data;
            2'b11   : regbus_A = ~{59'b0,rs1_r_addr};
            default : regbus_A = rs1_data;
        endcase
    end

    // Regbus B
    assign regbus_B = is_system_inst ?  csr_data  :  rs2_data  ;

    //----------Immediate Generate----------//
    wire [`XLEN-1:0] imm_I_shift;
    wire [`XLEN-1:0] imm_I_shift_w;
    wire [`XLEN-1:0] imm_I_nonshift;
    wire [`XLEN-1:0] imm_S;
    wire [`XLEN-1:0] imm_B;
    wire [`XLEN-1:0] imm_U;
    wire [`XLEN-1:0] imm_J;
    reg  [ 2: 0] imm_extension_type;

    assign imm_I_shift    = { 58'b0 , shamt };
    assign imm_I_shift_w  = { 59'b0 , shamt[4:0] };
    assign imm_I_nonshift = { {52{inst[31]}} , inst[31:20] };
    assign imm_B          = { {51{inst[31]}} , inst[31] , inst[7] , inst[30:25] , inst[11:8] , 1'b0 };
    assign imm_S          = { {52{inst[31]}} , inst[31:25] , inst[11:7] };
    assign imm_U          = { {32{inst[31]}} , inst[31:12] , 12'b0 };
    assign imm_J          = { {44{inst[31]}} , inst[19:12] , inst[20] , inst[30:21] , 1'b0 };

    always@(*) begin
        casez(opcode)
            5'b001?0 : imm_extension_type = (~funct3[1] & funct3[0] & ~opcode[1]) ?  `IMM_I_S_EXT    : 
                                            (~funct3[1] & funct3[0] & opcode[1])  ?  `IMM_I_S_W_EXT  :  `IMM_I_NS_EXT  ;
            5'b00000 : imm_extension_type = `IMM_I_NS_EXT;
            5'b11001 : imm_extension_type = `IMM_I_NS_EXT;
            `B_TYPE  : imm_extension_type = `IMM_B_EXT;
            `S_TYPE  : imm_extension_type = `IMM_S_EXT;
            `J_TYPE  : imm_extension_type = `IMM_J_EXT;
            `U_TYPE  : imm_extension_type = `IMM_U_EXT;
            default  : imm_extension_type = `IMM_NOEXT;
        endcase
    end

    always @(*) begin
        case (imm_extension_type)
            `IMM_I_S_EXT   : imm = imm_I_shift;
            `IMM_I_S_W_EXT : imm = imm_I_shift_w;
            `IMM_I_NS_EXT  : imm = imm_I_nonshift;
            `IMM_B_EXT     : imm = imm_B;
            `IMM_S_EXT     : imm = imm_S;
            `IMM_U_EXT     : imm = imm_U;
            `IMM_J_EXT     : imm = imm_J;
            default        : imm = 64'b0;
        endcase
    end

    //----------ALU----------//
    // ALU Oprend A
    /* 00:pc; 01:regbusA; 10:zext(bitsel[regbus_A]); 11:sext(bitsel[regbus_A]); */
    assign ALU_oprend_A_src[0] = ~((opcode[4] & opcode[3] & ~opcode[2] & opcode[0])                                          /* jal / jalr */
                               | (~opcode[4] & ~opcode[3] & opcode[2] & ~opcode[1] & opcode[0])                               /* auipc */
                               | (ALU_out_ext_type[0] & funct3[2] & ~funct3[1] & funct3[0] & ~funct7[5]));                     /* srlw / srliw */
    assign ALU_oprend_A_src[1] = ALU_out_ext_type[0] & funct3[2] & ~funct3[1] & funct3[0];                                 /* srlw / srliw / sraw / sraiw */
    assign rs1_r_en = (|ALU_oprend_A_src & ~opcode[0]) 
                    | (opcode[4] & opcode[3] & ~opcode[2] & ~opcode[1] & opcode[0]) ;
    
    // ALU Oprend B
    /* 00:imm; 01:const 4; 10:regbus_B; 11:zext(bitsel[regbus_B]); */
    assign ALU_oprend_B_src[0] = (ALU_out_ext_type[0] & opcode[3] & |funct3) 
                               | (opcode[4] & opcode[3] & ~opcode[2] & opcode[0])
                               | (~opcode[4] & opcode[3] & opcode[2] & ~opcode[1] & ~opcode[0] & ~funct3[1] & funct3[0]);
    assign ALU_oprend_B_src[1] = (~(is_system_inst & ~funct3[1] & funct3[0]) | ~is_system_inst) 
                               & ((ALU_out_ext_type[0] & opcode[3]) | ~opcode[4] & opcode[3] & opcode[2] & ~opcode[0] | opcode[4] & opcode[3] & (opcode[2] | ~opcode[2] & ~opcode[1] & ~opcode[0]));
    assign rs2_r_en = ALU_oprend_B_src[1] 
                    | (~opcode[4] & opcode[3] & ~opcode[2] & ~opcode[1] & ~opcode[0]);
    
    // ALU Operation
    always@(*)begin
        casez({opcode,funct3})
            {5'b110?1,3'b???}    : ALU_op = `ADD_OP;                             /* jal / jalr */
            {5'b0?101,3'b???}    : ALU_op = `ADD_OP;                             /* U_type */
            {5'b0?000,3'b???}    : ALU_op = `ADD_OP;                             /* load / store */
            {5'b001?0,3'b000}    : ALU_op = `ADD_OP;                             /* addi(w) */
            {`R_TYPE,3'b000}     : ALU_op = funct7[5] ? `SUB_OP  :  `ADD_OP  ;    /* R_add(w) / sub(w) */
            {`B_TYPE,3'b???}     : ALU_op = `SUB_OP;                             /* B_type */
            {5'b0?100,3'b01?}    : ALU_op = `SUB_OP;                             /* slti / sltiu / slt / sltu */
            {5'b0?100,3'b111}    : ALU_op = `AND_OP;
            {5'b0?100,3'b110}    : ALU_op = `OR_OP;
            {5'b0?100,3'b100}    : ALU_op = `XOR_OP;
            {5'b0?1?0,3'b001}    : ALU_op = `SLL_OP;
            {5'b0?1?0,3'b101}    : ALU_op = funct7[5] ? `SRA_OP  :  `SRL_OP  ;
            {`PRIVILEGED,3'b?10} : ALU_op = `OR_OP;
            {`PRIVILEGED,3'b?11} : ALU_op = `AND_OP;
            default              : ALU_op = `ADD_OP; 
        endcase
    end

    // ALU Output Extension
    /* [0]:64bit; [1]:cmp type;(1 unsigned / 0 signed) [2]:is_cmp */
    assign ALU_out_ext_type[0] = opcode[2] & opcode[1] & ~opcode[0];
    assign ALU_out_ext_type[1] = funct3[0];
    assign ALU_out_ext_type[2] = ~opcode[4] & opcode[2] & ~opcode[1] & ~opcode[0] & ~funct3[2] & funct3[1];

    //----------Branch/Jump Signals----------//
    assign pc_o = pc_i;
    assign pc_valid_o = pc_valid_i;
    assign inst_jump = opcode[4] & opcode[3] & ~opcode [2] & opcode [0];
    assign adder_src = opcode[4] & opcode[3] & ~opcode [2] & ~opcode[1] & opcode [0];
    assign inst_branch_type = (opcode[4] & opcode[3] & ~opcode[2] & ~opcode[1] & ~opcode[0]) ? funct3 : `NOTBRANCH ;

    //----------Privileged Instructions----------//
    assign is_system_inst = opcode[4] & opcode[3] & opcode[2] & ~opcode[1] & ~opcode[0];
    assign inst_ecall     = is_system_inst & (funct12 == 12'h000);
    assign inst_mret      = is_system_inst & (funct12 == 12'h302);
    assign inst_fencei    = (~opcode[4] & ~opcode[3] & ~opcode[2] & opcode[1] & opcode[0]);
    assign inst_wfi       = is_system_inst & ~(|funct3) & (funct12 == 12'b000100000101);

    //----------Memory/Regfile/CSR Access----------//
    // Memory
    assign mem_r_en = rd_w_en & rd_w_src;
    assign mem_r_ext_type = (rd_w_src & rd_w_en) ? funct3 : 3'b0 ;
    assign mem_w_en = ~opcode[4] & opcode[3] & ~opcode[2] & ~opcode[1] & ~opcode[0];
    assign mem_w_size = funct3[1:0];                                            /* 00:byte;  01:half word;  10:word;  11:double word; */

    // Regfile
    assign rd_w_en = ~opcode[3] | opcode[2] | opcode[1] | opcode[0];
    assign rd_w_src = ~(|opcode);                                          /* 1:alu 0:csr (csr_related) else 1:mem 0:alu */

    // CSR
    assign csr_r_en = is_system_inst;
    assign csr_w_en = is_system_inst & ~inst_mret & ~inst_ecall;
    
endmodule